import copy
import os
import random
import shutil
import subprocess
import math

import numpy as np
from seis2frac_funs import write_string


def compliklihood(fx_series, test, mole=1, unit=1, sigma=1):
    sum1 = sum2 = 0
    for i in range(len(fx_series)):  # error of simulation and observation
        e = (fx_series[i] - test[i]) ** 2
        sum1 += e * (mole * unit) ** 2  # unit change from mol/L to mg/L in liklihood computation
        sum2 += e
    rmse = np.sqrt(sum2 / len(fx_series))
    if sigma == 0:
        log_new = 0.0
    else:
        log_new = - sum1 / (sigma ** 2) / 2
    return rmse, log_new


def write_current_accept(cur, itera, p_acc, cons_2, string=''):  # write newly accepted values
    logs, rmse, dist, frac, k, cen, eqr = p_acc[0], p_acc[1], p_acc[2], p_acc[3], p_acc[4], p_acc[5], p_acc[6]
    fc, fs, fo = convert_frac_params(frac, cen, eqr)
    wrat = f'{dist},{k[0]}\n{fo}\n{fc}\n{fs}\n{rmse},{logs}'
    write_string('accept.txt', 'a', f'{itera}\t{rmse:.8f}\t{logs}\t{k[0]}\t{dist}\t{len(frac)}\n')
    write_string('current_fit.txt', 'w', wrat)
    stacc = f'RMSE={rmse * 100}%\n'
    for d in cons_2:  # write accepted simulation data
        stacc += f'{d}\n'
    desti = f'{cur}/accept/f_{itera}_{rmse:.5f}'
    shutil.copytree(f'{cur}/f', desti)
    write_string(f'{desti}/simulation_{itera}_{rmse:.5f}.txt', 'a', stacc)
    write_string('geometry_log.txt', 'a', f'{itera}\n{wrat}\n\n')
    if not string == '':
        write_string('inverse_log.txt', 'a', string)
    print(f'new comprehensive RMSE={rmse}')


def convert_frac_params(frac, cen, eqr):
    write_center, write_size, write_occurrence = '', '', ''
    for d in range(len(frac)):  # stored in type 'list'
        write_size += f'{eqr[d]:.5e}'
        for fg in range(3):
            write_center += f'{cen[d][fg]:.4e}'
            if fg < 2:
                write_center += ','
        if d < len(frac) - 1:
            write_size += ','
            write_center += ';'
        for e in range(5):
            write_occurrence += f'{frac[d][e]}'
            if e == 4:
                if d < len(frac) - 1:
                    write_occurrence += ';'  # separate different fractures
            else:
                write_occurrence += ','  # separate different parameters of a fracture
    return write_center, write_size, write_occurrence


def accept_or_reject(p_new, f_min, f_max, con_conv, m_acc, m_rej, p_old=None, iters=1):
    n_old, p_acc, l_a_r, Z = 0, copy.deepcopy(p_new), 0.0, 0.0
    cut = os.getcwd()
    if p_old is None:  # first successful simulation, directly accepted
        write_current_accept(cut, iters, p_acc, con_conv)
    else:
        n_old, move_type, log_l_h, log_pro = len(p_old[3]), 'update', p_new[0] - p_old[0], p_old[2] - p_new[2]
        if len(p_new[3]) > len(p_old[3]):  # assume maximum parameter dimension change is 1
            move_type = 'birth'
        elif len(p_new[3]) < len(p_old[3]):
            move_type = 'death'  # death move: using existed value for fix-dimension variable
        l_a_r, Z = min(log_l_h + log_pro, 0.0000), np.log(random.random())  # accept ratio
        print(f'iter {iters} {move_type}: log accept ratio={l_a_r:.4e}, alpha={Z:.4e}, accept={l_a_r > Z}')
        inv_res = f'{iters}\t{log_l_h:.4e}\t{log_pro:.4e}\t{p_new[1]}\t{len(p_new[3])}\t{l_a_r > Z}\n'
        if l_a_r > Z:
            write_current_accept(cut, iters, p_acc, con_conv, inv_res)
        else:
            p_acc = copy.deepcopy(p_old)
            print(f'present RMSE={p_acc[1]}, rejected RMSE={p_new[1]}')
            if p_new[1] / p_acc[1] <= 1.5:
                strej = f'RMSE={p_new[1] * 100}%\n'
                for d in con_conv:
                    strej += f'{d}\n'
                desti = f'{cut}/reject/f_{iters}_{p_new[1]:.5e}'
                shutil.copytree(f'{cut}/f', desti)
                write_string(f'{desti}/simulation_{iters}_{p_new[1]:.5e}.txt', 'a', strej)
                frc, frs, fro = convert_frac_params(p_new[3], p_new[5], p_new[6])
                geo_rea = f'{iters}\n{p_new[2]},{p_new[4][0]}\n{fro}\n{frc}\n{frs}\n{p_new[1]},{p_new[0]}\n\n'
                write_string('geometry_log.txt', 'a', geo_rea)
                write_string('inverse_log.txt', 'a', inv_res)
    next_f = boundary_handling(len(p_acc[3]), f_min, f_max, m_acc, m_rej, n_old, iters, l_a_r > Z)
    return p_acc, next_f


def boundary_handling(num_new, fmin, fmax, p_a, p_r, num_old=0, iters=1, accept=True):
    q = random.random()
    move = 'update'
    if num_old == 0:  # first simulation, assign equal probability for each possible move
        if (num_new - fmin) * (num_new - fmax) == 0:  # number of fractures equal to boundary
            if q <= 0.5:
                if num_new == fmin:
                    num_new, move = num_new + 1, 'birth'
                elif num_new == fmax:
                    num_new, move = num_new - 1, 'death'
        else:  # three possible moves
            if 0 <= q * 3 < 1:
                num_new, move = num_new + 1, 'birth'
            elif 1 <= q * 3 < 2:
                num_new, move = num_new - 1, 'death'
    else:
        if num_new > num_old:  # birth move applied
            if accept:  # probability:[0.6, 0.2, 0.2], num_new is accepted
                if num_new < fmax:  # number of fractures does not reach the maximum, 3 possible moves
                    if q < p_a:
                        num_new, move = num_new + 1, 'birth'
                    elif q < (0.5 + p_a / 2):
                        num_new, move = num_new - 1, 'death'
                elif q < 0.5:  # only death and update moves with equal probability
                    num_new, move = num_new - 1, 'death'
            else:  # probability:[0.1, 0.45, 0.45], num_old is accepted
                if num_old == fmin:
                    if q < (2 * p_r / (1 + p_r)):
                        num_new, move = num_old + 1, 'birth'
                    else:
                        num_new = num_old
                else:
                    if q < p_r:
                        num_new, move = num_old + 1, 'birth'
                    elif q < (0.5 + p_r / 2):
                        num_new, move = num_old - 1, 'death'
                    else:
                        num_new = num_old
        elif num_new < num_old:  # death move applied
            if accept:  # probability:[0.2, 0.6, 0.2]
                if num_new > fmin:  # number of fractures does not reach the minimum, 3 possible moves
                    if q < (0.5 - p_a / 2):
                        num_new, move = num_new + 1, 'birth'
                    elif q < (0.5 + p_a / 2):
                        num_new, move = num_new - 1, 'death'
                elif q < 0.5:  # only birth and update moves with equal probability
                    num_new, move = num_new + 1, 'birth'
            else:  # probability:[0.45, 0.1, 0.45]
                if num_old == fmax:
                    if q < (2 * p_r / (1 + p_r)):
                        num_new, move = num_old - 1, 'death'
                    else:
                        num_new = num_old
                else:
                    if q < (0.5 - p_r / 2):
                        num_new, move = num_old + 1, 'birth'
                    elif q < (0.5 + p_r / 2):
                        num_new, move = num_old - 1, 'death'
                    else:
                        num_new = num_old
        else:  # update move applied
            if accept:  # probability:[0.2, 0.2, 0.6]
                if num_new == fmin:  # only birth and update moves
                    if q < ((1 - p_a) / (1 + p_a)):
                        num_new, move = num_new + 1, 'birth'
                elif num_new == fmax:  # only death and update moves
                    if q < ((1 - p_a) / (1 + p_a)):
                        num_new, move = num_new - 1, 'death'
                else:
                    if q < ((1 - p_a) / 2):
                        num_new, move = num_new + 1, 'birth'
                    elif q < (1 - p_a):
                        num_new, move = num_new - 1, 'death'
            else:  # probability:[0.45, 0.45, 0.1]
                if num_old == fmin:  # only birth and update moves
                    if q < ((1 - p_r) / (1 + p_r)):
                        num_new, move = num_old + 1, 'birth'
                    else:
                        num_new = num_old
                elif num_old == fmax:  # only death and update moves
                    if q < ((1 - p_r) / (1 + p_r)):
                        num_new, move = num_old - 1, 'death'
                    else:
                        num_new = num_old
                else:
                    if q < ((1 - p_r) / 2):
                        num_new, move = num_old + 1, 'birth'
                    elif q < (1 - p_r):
                        num_new, move = num_old - 1, 'death'
                    else:
                        num_new = num_old
    print(f'random={q}\tnext={num_new}\tmove={move}')
    write_string('next.txt', 'a', f'next_move:iter={iters}\tnext={num_new}\tmove={move}\t')
    return num_new

